package lte

import (
	"math/cmplx"
	"sort"

	"gitee.com/thubcc/mobile/utils"
)

var (
	offset     = []int{10, 558, 970, 1518}
	symbol_idx = []int{0, 4, 7, 11}
)

type LazyReferences struct {
	ref map[int]*RefSignal
}

type RefSignal struct {
	CRS
	ref    utils.Array3
	lazyFd map[int][]complex128
}

func BuildRef(cid, nof_ports, nof_prb int, cp string) *RefSignal {
	crs := NewCRS(cid, nof_ports, nof_prb, cp)
	refd := make([][][]complex128, nof_ports)
	for p := range refd {
		refd[p] = crs.cs_tx(p, false, false).D
	}
	return &RefSignal{
		CRS:    crs,
		ref:    utils.Array3dByD(refd),
		lazyFd: make(map[int][]complex128),
	}
}

func NewLazyRef() *LazyReferences {
	return &LazyReferences{
		make(map[int]*RefSignal),
	}
}

func (l *LazyReferences) Get(cid int) *RefSignal {
	if r, ok := l.ref[cid]; ok {
		return r
	} else {
		r = BuildRef(cid, 4, 6, "normal")
		l.ref[cid] = r
		return r
	}
}

func (self *RefSignal) frame_refsignal(port int, buf []complex128) {
	for slot, v := range self.ref.D[port] {
		o := slot * 1920
		copy(buf[o:o+1920], v)
	}
	return
}

func (self *RefSignal) GetFd(port int) (r []complex128) {
	var ok bool
	if r, ok = self.lazyFd[port]; ok {
		return
	} else {
		td2fft, err := td2fd.Get()
		if err != nil {
			return
		}
		defer td2fd.Return(td2fft)
		self.frame_refsignal(port, td2fft.Buf)
		td2fft.Execute()
		r = make([]complex128, len(td2fft.Buf))
		copy(r, td2fft.Buf)
		self.lazyFd[port] = r
	}
	return
}

func (self *RefSignal) SymOfSlot(buf []complex128, port, slot, peak int) (ch [4]complex128) {
	symbol_sz := LTE_SYMBOL_SZ[self.NofPrb]
	for i, off := range offset {
		d := buf[off+peak : off+peak+symbol_sz]
		r := self.ref.D[port][slot][off : off+symbol_sz]
		for j, v := range d {
			ch[i] += v * cmplx.Conj(r[j])
		}
	}
	return
}

func (self *RefSignal) SymOfFrame(buf []complex128, port, peak int, ch []complex128) {
	symbol_sz := LTE_SYMBOL_SZ[self.NofPrb]
	sf_len := LTE_SF_LEN(symbol_sz)
	sf_idx_init := (2*LTE_NOF_SF_X_FRAME - peak/sf_len) % LTE_NOF_SF_X_FRAME
	n := peak % sf_len
	for _slot := 0; _slot < LTE_NOF_SF_X_FRAME; _slot++ {
		slot := (sf_idx_init + _slot) % LTE_NOF_SF_X_FRAME
		sym := self.SymOfSlot(buf, port, slot, n)
		for i, v := range sym[:] {
			ch[_slot*sf_len+offset[i]] = v
		}
		n += sf_len
	}
}

func (self *RefSignal) XcorrSlot(buf []complex128, port, slot, peak int, ch []complex128) {
	symbol_sz := LTE_SYMBOL_SZ[self.NofPrb]
	sf_len := LTE_SF_LEN(symbol_sz)
	r := self.ref.D[port][slot]
	for i := -64; i < 64; i++ {
		if peak+i >= 0 {
			d := buf[peak+i : peak+i+sf_len]
			vr, _ := utils.NpDot(d, r)
			ch[i+64] += vr
		} else {
			d := buf[:sf_len+peak+i]
			vr, _ := utils.NpDot(d, r[-i-peak:sf_len])
			ch[i+64] += vr
		}
	}
	return
}

func (self *RefSignal) H(buf []complex128, port, peak int) (ch [128]complex128) {
	symbol_sz := LTE_SYMBOL_SZ[self.NofPrb]
	sf_len := LTE_SF_LEN(symbol_sz)
	sf_idx_init := (2*LTE_NOF_SF_X_FRAME - peak/sf_len) % LTE_NOF_SF_X_FRAME
	n := peak % sf_len
	for _slot := 0; _slot < LTE_NOF_SF_X_FRAME; _slot++ {
		slot := (sf_idx_init + _slot) % LTE_NOF_SF_X_FRAME
		self.XcorrSlot(buf, port, slot, n, ch[:])
		n += sf_len
	}
	return
}

func (self *RefSignal) Pos0(buf []complex128, port, peak int) []complex128 {
	res := make([]complex128, LTE_NOF_SF_X_FRAME)
	symbol_sz := LTE_SYMBOL_SZ[self.NofPrb]
	sf_len := LTE_SF_LEN(symbol_sz)
	sf_idx_init := (2*LTE_NOF_SF_X_FRAME - peak/sf_len) % LTE_NOF_SF_X_FRAME
	n := peak % sf_len
	for _slot := 0; _slot < LTE_NOF_SF_X_FRAME; _slot++ {
		slot := (sf_idx_init + _slot) % LTE_NOF_SF_X_FRAME
		for _, off := range offset {
			r, _ := utils.NpDot(buf[off+n:off+n+symbol_sz], self.ref.D[port][slot][off:off+symbol_sz])
			res[_slot] += r
		}
		n += sf_len
	}
	return res
}

func (self *RefSignal) AFCSolt(buf []complex128, port, peak int) float64 {
	l := len(buf)
	p := (l - peak) / 19200
	r := make([][]complex128, p-1)
	for i := range r {
		r[i] = self.Pos0(buf[i*19200:i*19200+19200], port, peak)
	}
	fft := utils.NewFFT((p - 1) * LTE_NOF_SF_X_FRAME)
	for i, v := range r {
		copy(fft.Buf[i*LTE_NOF_SF_X_FRAME:i*LTE_NOF_SF_X_FRAME+LTE_NOF_SF_X_FRAME], v)
	}
	fft.Execute()
	fftl := (p - 1) * LTE_NOF_SF_X_FRAME
	idx, _ := utils.CAbsMax1D(fft.Buf)
	if idx > fftl/2 {
		idx -= fftl
	}
	return float64(idx*1000) / float64(fftl)
}

var (
	lazy = NewLazyRef()
)

func AFCSoltofCell(buf []complex128, cid, port, peak int) float64 {
	rs := lazy.Get(cid)
	return rs.AFCSolt(buf, port, peak)
}

func HofCell(buf []complex128, cid, port, peak int) [128]complex128 {
	rs := lazy.Get(cid)
	return rs.H(buf, port, peak)
}

type ScanRefResult struct {
	Cid int
	Pos int
	Df  float64
	Amp float64
}

type ScanRefResults []*ScanRefResult

func (rs ScanRefResults) Len() int           { return len(rs) }
func (rs ScanRefResults) Less(i, j int) bool { return rs[i].Amp > rs[j].Amp }
func (rs ScanRefResults) Swap(i, j int)      { rs[i], rs[j] = rs[j], rs[i] }

func ScanRefSym(buf []complex128, t, port, peak int) (rs ScanRefResults) {
	rs = make(ScanRefResults, 0)
	for cid := t; cid < 504; cid += 3 {
		rs = append(rs, RefAFC(buf, cid, port, peak))
		rs = append(rs, RefAFC(buf, cid, port, (peak+9600)%19200))
	}
	sort.Sort(rs)
	return
}

func RefAFC(buf []complex128, cid, port, peak int) *ScanRefResult {
	l := 19200 * 10
	fft := utils.NewFFT(l)
	rs := lazy.Get(cid)
	rs.SymOfFrame(buf, port, peak, fft.Buf[:])
	fft.Execute()
	pidx, pamp := utils.CAbsMax1D(fft.Buf[:500])
	nidx, namp := utils.CAbsMax1D(fft.Buf[l-500:])
	idx := 0
	var amp float64
	if pamp > namp {
		idx = pidx
		amp = pamp
	} else {
		idx = nidx - 500
		amp = namp
	}
	df := float64(idx * 10)
	pos := peak
	return &ScanRefResult{cid, pos, df, amp}
}

func (l *LazyReferences) FindPeak(buf []complex128, cid, port int) (idx int, amp float64) {
	len0 := len(buf)
	p := len0 % 19200
	s := len0 - p - 19200
	e := len0 - p
	ref := l.Get(cid)
	fd := ref.GetFd(port)
	fft, err := td2fd.Get()
	if err != nil {
		return
	}
	defer td2fd.Return(fft)
	copy(fft.Buf, buf[s:e])
	fft.Execute()
	ifft, ierr := xcorrIFFT.Get()
	if ierr != nil {
		return
	}
	defer xcorrIFFT.Return(ifft)
	for i, v := range fd {
		ifft.Buf[i] = fft.Buf[i] * cmplx.Conj(v)
	}
	ifft.Execute()
	idx, amp = utils.CAbsMax1D(ifft.Buf)
	return
}

func (l *LazyReferences) ScanFreq(buf []complex128, cid, port int, frange []int) (idx, df int, amp float64) {
	ref := l.Get(cid)
	fd := ref.GetFd(port)
	flen := len(frange)
	famp := make([]float64, flen)
	fidx := make([]int, flen)
	fftlen := len(fd)
	fft, err := td2fd.Get()
	if err != nil {
		return
	}
	defer td2fd.Return(fft)
	ifft, ierr := xcorrIFFT.Get()
	if ierr != nil {
		return
	}
	defer xcorrIFFT.Return(ifft)

	copy(fft.Buf, buf[:19200])
	fft.Execute()
	for fi, s := range frange {
		for i, v := range fft.Buf {
			ifft.Buf[i] = v * cmplx.Conj(fd[(i-s+fftlen)%fftlen])
		}
		ifft.Execute()
		fidx[fi], famp[fi] = utils.CAbsMax1D(ifft.Buf)
	}
	idx, amp = utils.AbsMax1D(famp)
	df = frange[fidx[idx]]
	return
}

func RefSignalPeak(buf []complex128, cid, port int) (int, float64) {
	return lazy.FindPeak(buf, cid, port)
}

func RefSignalH(buf []complex128, cid, port, pos int) [128]complex128 {
	rs := lazy.Get(cid)
	return rs.H(buf, port, pos)
}

func ScanSignalPeak(buf []complex128, port int) (rs ScanRefResults) {
	rs = make(ScanRefResults, 0)
	for id := 0; id < 504; id++ {
		pos, amp := RefSignalPeak(buf, id, port)
		rs = append(rs, &ScanRefResult{id, pos, 0.0, amp})
	}
	sort.Sort(rs)
	return
}
